#!/usr/bin/env python3

import os
import shutil

import common
from shell_helpers import LF

class Main(common.BuildCliFunction):
    def __init__(self):
        super().__init__(
            description='''\
Build the Linux kernel.
'''
        )
        self.add_argument(
            '--config', default=[], action='append',
            help='''\
Add a single kernel config configs to the current build. Sample value:
'CONFIG_FORTIFY_SOURCE=y'. Can be used multiple times to add multiple
configs. Takes precedence over any config files.
'''
        )
        self.add_argument(
            '--config-fragment', default=[], action='append',
            help='''\
Also use the given kernel configuration fragment file.
Pass multiple times to use multiple fragment files.
'''
        )
        self.add_argument(
            '--build',
            default=True,
            help='''\
Build the kernel.
'''
        )
        self.add_argument(
            '--configure',
            default=True,
            help='''\
Configure the kernel.
'''
        )
        self.add_argument(
            '--custom-config-file',
            help='''\
Use this file as the .config. Don't add any default framents to it,
unless explicitly passed with `--config` and `--config-fragment` on
top of it.
'''
        )
        self.add_argument(
            '--custom-config-file-gem5',
            default=False,
            help='''\
Like --custom-config-file, but select the gem5 Linux kernel fork
config as the custom config file. Ignore --custom-config-file if given.
See: https://cirosantilli.com/linux-kernel-module-cheat#gem5-arm-linux-kernel-patches
'''
        )
        self.add_argument(
            '--custom-config-target',
            help='''\
Like --custom-config-file, but generate the base configuration file
by running a kernel make target such as menuconfig or defconfig.
If a .config exists in the tree, it will get picked e.g. by menuconfig,
so you might want to --clean the build first.
'''
        )
        self.add_argument(
            '--modules-install',
            default=True,
            help='''\
Run `make modules_install` after `make`.
'''
        )
        self.add_argument(
            'extra_make_args',
            default=[],
            metavar='extra-make-args',
            nargs='*'
        )
        self._add_argument('--force-rebuild')

    def build(self):
        build_dir = self.get_build_dir()
        os.makedirs(build_dir, exist_ok=True)
        common_args = {
            'cwd': self.env['linux_source_dir'],
        }
        ccache = shutil.which('ccache')
        if ccache is not None:
            cc = '{} {}'.format(ccache, self.env['gcc_path'])
        else:
            cc = self.env['gcc_path']
        if self.env['verbose']:
            verbose = ['V=1']
        else:
            verbose = []
        common_make_args = [
            'make', LF,
            '-j', str(self.env['nproc']), LF,
            'ARCH={}'.format(self.env['linux_arch']), LF,
            'CROSS_COMPILE={}'.format(self.env['toolchain_prefix_dash']), LF,
            'CC={}'.format(cc), LF,
            'O={}'.format(build_dir), LF,
        ] + verbose
        if self.env['force_rebuild']:
            common_make_args.extend(['-B', LF])
        if self.env['configure']:
            if self.env['custom_config_target']:
                base_config_given = True
                base_config_needs_copy = False
            elif self.env['custom_config_file_gem5']:
                base_config_given = True
                base_config_needs_copy = True
                custom_config_file = os.path.join(
                    self.env['linux_source_dir'],
                    'arch',
                    self.env['linux_arch'],
                    'configs',
                    'gem5_defconfig'
                )
            elif self.env['custom_config_file']:
                base_config_given = True
                base_config_needs_copy = True
                custom_config_file = self.env['custom_config_file']
            else:
                base_config_given = False
                base_config_needs_copy = True
            if base_config_given:
                if base_config_needs_copy:
                    if not os.path.exists(custom_config_file):
                        raise Exception('config fragment file does not exist: {}'.format(custom_config_file))
                    base_config_file = custom_config_file
                config_fragments = []
            else:
                base_config_file = os.path.join(
                    self.env['linux_config_dir'],
                    'buildroot-{}'.format(self.env['arch'])
                )
                config_fragments = ['min', 'default']
                for i, config_fragment in enumerate(config_fragments):
                    config_fragments[i] = os.path.join(
                        self.env['linux_config_dir'],
                        config_fragment
                    )
            config_fragments.extend(self.env['config_fragment'])
            cli_configs = self.env['config']
            if self.env['initramfs']:
                cli_configs.append('CONFIG_INITRAMFS_SOURCE="{}"'.format(self.env['buildroot_cpio']))
            if cli_configs:
                cli_config_fragment_path = os.path.join(build_dir, 'lkmc_cli_config_fragment')
                self.sh.write_configs(cli_config_fragment_path, cli_configs, mode='w')
                config_fragments.append(cli_config_fragment_path)
            if base_config_needs_copy:
                self.sh.cp(
                    base_config_file,
                    os.path.join(self.env['linux_config']),
                )
            if self.env['custom_config_target']:
                self.sh.run_cmd(
                    (
                        common_make_args +
                        [self.env['custom_config_target'], LF]
                    ),
                    **common_args
                )
            if config_fragments:
                self.sh.run_cmd(
                    [
                        os.path.join(
                            self.env['linux_source_dir'],
                            'scripts',
                            'kconfig',
                            'merge_config.sh'
                        ), LF,
                        '-m', LF,
                        '-O', build_dir, LF,
                        os.path.join(self.env['linux_config']), LF,
                    ] +
                    self.sh.add_newlines(config_fragments)
                )
            self.sh.run_cmd(
                (
                    common_make_args +
                    ['olddefconfig', LF]
                ),
                **common_args
            )
        if self.env['build']:
            self.sh.run_cmd(
                (
                    common_make_args +
                    self.sh.add_newlines(self.env['extra_make_args'])
                ),
                # https://cirosantilli.com/linux-kernel-module-cheat#proc-version
                extra_env={
                    'KBUILD_BUILD_VERSION': '1',
                    'KBUILD_BUILD_TIMESTAMP': 'Thu Jan  1 00:00:00 UTC 1970',
                    'KBUILD_BUILD_USER': self.env['repo_short_id'],
                    'KBUILD_BUILD_HOST': common.git_sha(self.env['linux_source_dir']),
                },
                **common_args
            )
            if self.env['modules_install']:
                self.sh.run_cmd(
                    (
                        common_make_args +
                        [
                            'INSTALL_MOD_PATH={}'.format(self.env['out_rootfs_overlay_lkmc_dir']), LF,
                            'modules_install', LF,
                        ]
                    ),
                    **common_args
                )
                # TODO: remove build and source https://stackoverflow.com/questions/13578618/what-does-build-and-source-link-do-in-lib-modules-kernel-version
                # TODO Basically all kernel modules also basically leak full host paths. Just terrible. Buildroot deals with that stuff nicely for us.
                # self.rmrf()

    def get_build_dir(self):
        return self.env['linux_build_dir']

if __name__ == '__main__':
    Main().cli()
